package com.treasure.infrastructure.index.impl;

import com.alibaba.dashscope.exception.NoApiKeyException;
import com.treasure.common.exception.ServiceException;
import com.treasure.infrastructure.index.IndexDocHandler;
import com.treasure.infrastructure.index.IndexDocTypeEnum;
import com.treasure.infrastructure.index.loader.ESSegmentLoader;
import com.treasure.infrastructure.model.FileSegmentBean;
import com.treasure.common.model.UploadDocumentRequest;
import com.treasure.infrastructure.model.doc.SegmentVectorDocBean;
import com.treasure.infrastructure.repository.extraction.vector.SegmentVectorDocRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;


@Service("indexVectorDoc")
@RequiredArgsConstructor
@Slf4j
public class IndexVectorDocHandlerImpl extends ESSegmentLoader implements IndexDocHandler {

    private final SegmentVectorDocRepository segmentVectorDocRepository;

    @Override
    public IndexDocTypeEnum getIndexDocType() {
        return IndexDocTypeEnum.INDEX_VECTOR_DOC;
    }

    @Override
    public void indexResourceSegment(String organizationCode, List<String> resourcePaths) throws IOException, NoApiKeyException {
        List<FileSegmentBean> resourceSegmentBeans = loadSegmentList(resourcePaths);
        UploadDocumentRequest uploadDocumentRequest = new UploadDocumentRequest();
        uploadDocumentRequest.setCorpId("");
        uploadDocumentRequest.setUserId("");
        uploadDocumentRequest.setKnowledgeId(0);
        uploadDocumentRequest.setDocumentId(1L);
        uploadDocumentRequest.setDocumentName("undefined");
        indexResourceSegment(uploadDocumentRequest, resourceSegmentBeans);
    }

    @Override
    public void indexResourceSegment(UploadDocumentRequest uploadDocumentRequest, List<FileSegmentBean> fileSegmentBeanList) throws NoApiKeyException {
        List<SegmentVectorDocBean> indexSegmentDocList = new ArrayList<>(fileSegmentBeanList.size() * 5);
        for (FileSegmentBean resourceSegmentBean : fileSegmentBeanList) {
            String segmentFirst = resourceSegmentBean.getSegmentList().stream().findFirst().orElse(null);
            if (Objects.isNull(segmentFirst)) continue;
            for (String segmentText : resourceSegmentBean.getSegmentList()) {
                SegmentVectorDocBean segmentDocBean = new SegmentVectorDocBean();
                segmentDocBean.setDocName(resourceSegmentBean.getFileName());
                segmentDocBean.setPassages(segmentText);
                indexSegmentDocList.add(segmentDocBean);
            }
        }

        if (CollectionUtils.isEmpty(indexSegmentDocList)) {
            log.error("indexResourceSegment error indexSegmentDocList is empty method return fileSegmentBeanList:{}", fileSegmentBeanList);
            throw new ServiceException("indexResourceSegment error indexSegmentDocList is empty method return");
        }
        segmentVectorDocRepository.indexVector(indexSegmentDocList, uploadDocumentRequest);
    }

}
